
#pragma once
#include <NvInfer.h>
#include <opencv2/opencv.hpp>
#include <string>

namespace mono_depth {
class DepthAnythingModel {
 public:
  DepthAnythingModel(const std::string& trt_engine_path);
  ~DepthAnythingModel();

  DepthAnythingModel(const DepthAnythingModel&) = delete;
  DepthAnythingModel& operator=(const DepthAnythingModel&) = delete;
  DepthAnythingModel(DepthAnythingModel&&) = delete;
  DepthAnythingModel& operator=(DepthAnythingModel&&) = delete;

  cv::Mat process(const cv::Mat input);

 private:
  void show_engine_bindings() const;
  void fill_input_tensor(const cv::Mat& input);

 private:
  class TrtLogger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override;
  };

  std::unique_ptr<nvinfer1::IRuntime> runtime_;
  std::unique_ptr<nvinfer1::ICudaEngine> engine_;
  std::unique_ptr<nvinfer1::IExecutionContext> context_;
  TrtLogger logger_;

  // input - output
  void* binding_buffers_[2] = {nullptr, nullptr};
  cudaStream_t stream_;

  // intermediate data
  cv::Mat scaled_input_;
  std::vector<cv::Mat> channels_;
  std::vector<cv::Mat> channels_f_;
  std::vector<float> input_tensor_;
  std::vector<int> depth_data_;
  cv::Mat output_image_;
  cv::Mat colormap_;
};
}  // namespace mono_depth
